from models.YOLOv1 import YOLOv1_resnet
from models.YOLOv1 import YOLOv1_Loss
from data.DataSet import VOC2012
from torch.utils.data import DataLoader
import torchnet
import numpy as np
import torch
import config
import visdom

if __name__ == "__main__":
    # 可选选项
    opt = config.DefaultConfig()
    # 训练集
    train_data = VOC2012()
    train_dataloader = DataLoader(VOC2012(is_train=True), batch_size=3, shuffle=True)
    # 模型
    model = YOLOv1_resnet().cuda()

    # 冻结预训练的模型参数
    for layer in model.children():
        layer.requires_grad = False
        break
    # 损失函数
    criterion = YOLOv1_Loss()
    optimizer = torch.optim.SGD(model.parameters(), lr=opt.lr, momentum=0.9, weight_decay=0.0005)
    #统计指标
    loss_meter = torchnet.meter.AverageValueMeter()

    for epoch_num in range(opt.epoch):
        model.train()
        yl = torch.tensor([0]).cuda()
        loss_meter.reset()
        for i, (inputs, labels) in enumerate(train_dataloader):
            inputs = inputs.cuda()
            labels = labels.float().cuda()
            pre = model(inputs)
            loss = criterion(pre, labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            loss_meter.add(loss.data.cpu())
            print("Epoch %d/%d| Step %d/%d| Loss: %.2f  | Avg_Loss:%.2f" % (epoch_num, opt.epoch, i, len(train_data) // opt.batchsize, loss, loss_meter.mean))
        if (epoch_num+1) % opt.save_num == 0:
            torch.save(model, "./checkpoint/epoch"+str(epoch_num+1)+".pkl")






